#coding=utf-8
from __future__ import division
import re
import mxnet as mx
import os, urllib ,sys
import cv2
import numpy as np
import pandas as pd
import logging
import glob
sys.path.append('../utils/')
from mxnet_helper import  get_image,gen_prediction2,label_mapping,gen_prediction


img_sz = 360
crop_sz = 320
batchsize = 256

# ###########################
# # loading finetuned model #
# ###########################
modelname = 'resnext-101'
modelmode = 'finetune'
modelfolder = '/media/hszc/model/zhangchi/BaiduImage_model/'
modedir = os.path.join(modelfolder,modelname ,'finetuned_model')
train_lst_path = '../data/trainall_train.lst'
epoch = 23
test_data_shape = (batchsize,3,320,320)

print('loading model from {0}'.format(modedir))
sym, arg_params, aux_params = mx.model.load_checkpoint('{0}/{1}'.format(modedir,modelname), epoch)
mod = mx.mod.Module(symbol=sym, context=mx.gpu(1), label_names=None)

mod.bind(for_training=False, data_shapes=[('data', test_data_shape)],
         label_shapes=mod._label_shapes)
mod.set_params(arg_params, aux_params, allow_missing=True)

# ###########
# # predict #
# ###########
imgfolder = '/media/hszc/data/BaiduImage/dataset/test1'
img_path_list =[]
imgnames = []
for idx,imgfile in enumerate(os.listdir(imgfolder)):
    imgnames.append(imgfile.split('.')[0])
    imgpath = os.path.join(imgfolder,imgfile)
    img_path_list.append(imgpath)

import itertools
flip_list = [True,False]
crop_list = ['center','top-left','top-right','bottom-left','bottom-right']
proba_mat = []

for  pair in list(itertools.product(flip_list,crop_list)):
    isflip = pair[0]
    crop_type = pair[1]
    print('isfilp:{0},crop_type:{1}'.format(isflip,crop_type))

    _,proba  = gen_prediction2(mod,img_path_list,batchsize,img_sz,crop_sz,isflip,crop_type)

    proba_mat.append(proba)
preds = np.array(proba_mat)   #10 10593 100
print preds.shape
print np.sum(preds,axis=0).shape
preds  = np.argmax(np.sum(preds,axis=0),axis=1)
print preds.shape
result = label_mapping(preds, imgnames,train_lst_path)

proba_mat = np.mean(np.array(proba_mat),axis=0)
proba_mat_tfidx = np.zeros(proba_mat.shape, dtype=float)

# label map
mx_label_map = pd.read_csv(train_lst_path,sep='\t',header=None,names = ['idx','mx_label','imgfile'])
mx_label_map['true_label'] = mx_label_map['imgfile'].str.split('/').str[0]
mx_label_map = mx_label_map[['true_label', 'mx_label']].drop_duplicates('true_label')
mx_label_map['true_label'] = mx_label_map['true_label'].astype(int)
tf_label_map = pd.read_csv('../disc.txt',skiprows=1,header=None,names = ['tf_label','true_label'],sep=' ')
label_map = pd.merge(tf_label_map,mx_label_map,on='true_label',how='left')

# 按照tf_label_map 重排列
for idx in label_map.index:
    tf_label = int(label_map.loc[idx,'tf_label'])
    mx_label = int(label_map.loc[idx,'mx_label'])
    proba_mat_tfidx[:, tf_label] = proba_mat[:, mx_label]  #按照tf_label顺序重排mx 的score


#重排行
tf_img_idx = pd.DataFrame({'img_path':sorted(glob.glob(imgfolder+'/*.jpg'))})


mx_score = pd.DataFrame(proba_mat_tfidx)
mx_score['img_path'] = img_path_list
proba_mat_tfidx = pd.merge(tf_img_idx, mx_score, on='img_path', how='left')
proba_mat_tfidx = proba_mat_tfidx.drop(['img_path'], axis=1).values

print proba_mat_tfidx.shape
np.save('results/{0}_{1}_{2}_flipcrop_merged_proba_mat_tfidx.npy'.format(modelname,modelmode,epoch),proba_mat_tfidx)


result = result.sort_values(by=['imgid'])
result[['label','imgid']].to_csv('results/{0}_{1}_{2}_flipcrop_merged.txt'.format(modelname,modelmode,epoch),
                                     index=False,sep='\t',header=None)




